from __future__ import absolute_import
from .resnet import *
from .selfpoolingdir import SelfPoolingDir
from .crosspoolingdir import CrossPoolingDir
from .classifier import Classifier
from .attmodel import AttModuleDir
from .classifier import Classifier


__factory = {
    'resnet18': resnet18,
    'resnet34': resnet34,
    'resnet50': resnet50,
    'resnet101': resnet101,
    'resnet152': resnet152,
    'attmodel': AttModuleDir,
    'classifier': Classifier,
}


def names():
    return sorted(__factory.keys())


def create(name, *args, **kwargs):
    """
        Create a model instance.
        Parameters
        ----------
        name : str
            Model name. Can be one of 'inception', 'resnet18', 'resnet34',
            'resnet50', 'resnet101', and 'resnet152'.
        pretrained : bool, optional
            Only applied for 'resnet*' models. If True, will use ImageNet pretrained
            model. Default: True
        cut_at_pooling : bool, optional
            If True, will cut the model before the last global pooling layer and
            ignore the remaining kwargs. Default: False
        num_features : int, optional
            If positive, will append a Linear layer after the global pooling layer,
            with this number of output units, followed by a BatchNorm layer.
            Otherwise these layers will not be appended. Default: 256 for
            'inception', 0 for 'resnet*'
        norm : bool, optional
            If True, will normalize the feature to be unit L2-norm for each sample.
            Otherwise will append a ReLU layer after the above Linear layer if
            num_features > 0. Default: False
        dropout : float, optional
            If positive, will append a Dropout layer with this dropout rate.
            Default: 0
        """
    if name not in __factory:
        raise KeyError("Unknown model:", name)
    return __factory[name](*args, **kwargs)
